/**
 * 
 */
package edu.umd.clip.lm.util;


public class CountDistribution {
	private Long2IntMap counts;
	private long totalCount;
	/**
	 * @param counts
	 * @param totalCount
	 */
	public CountDistribution(Long2IntMap counts, long totalCount) {
		this.counts = counts;
		this.totalCount = totalCount;
	}
	
	public CountDistribution() {
		this.counts = new Long2IntMap();
		this.totalCount = 0;
	}
	
	public Long2IntMap getCounts() {
		return counts;
	}
	public void setCounts(Long2IntMap counts) {
		this.counts = counts;
	}
	public long getTotalCount() {
		return totalCount;
	}
	public void setTotalCount(long totalCount) {
		this.totalCount = totalCount;
	}
	
	public static Long2IntMap pruneCounts(Long2IntMap counts, int minValue) {
		Long2IntMap newCounts = new Long2IntMap(counts.size());
		for(Long2IntMap.Iterator it = counts.iterator(); it.hasNext(); ) {
			Long2IntMap.Entry e = it.next();
			if (e.getValue() >= minValue) {
				newCounts.put(e.getKey(), e.getValue());
			}
		}
		return newCounts;
	}
	
	public void addCounts(Long2IntMap counts) {
		long addedCounts = 0;
		for(Long2IntMap.Iterator i = counts.iterator(); i.hasNext(); ) {
			Long2IntMap.Entry e = i.next();
			int count = e.getValue();
			addedCounts += count;
			this.counts.addAndGet(e.getKey(), count);
		}
		totalCount += addedCounts;
	}
	
	public void addDistribution(CountDistribution dist) {
		addCounts(dist.counts);
	}
	
	public void substractCounts(Long2IntMap counts) {
		long substractedCounts = 0;
		for(Long2IntMap.Iterator i = counts.iterator(); i.hasNext(); ) {
			Long2IntMap.Entry e = i.next();
			int count = e.getValue();
			substractedCounts += count;
			int newCount = this.counts.addAndGet(e.getKey(), -count);
			assert(newCount >= 0);
		}
		totalCount -= substractedCounts;
	}
	
	public void substractDistribution(CountDistribution dist) {
		substractCounts(dist.counts);
	}
}